# !/usr/bin/python3
# -*- coding: utf-8 -*-
# @Time : 2021/8/31 10:24
# @Author : Yu LiXinQian
# @Email : yulixinqian805@gmail.com
# @File : SVD.py
# @Project : My_MLF

import numpy as np


class SVD:

    def __init__(self, A):
        self.A = A
        self.U = None
        self.V = None
        self.sigmas = None

    def SVD(self, A, Num_Of_sigmas = None):
        # Compute the eigenvalues of A.T*A
        AT_A = np.dot(A.T, A)

        # Compute the eigenvalues and right eigenvectors of a square array
        w, V = np.linalg.eig(AT_A)  # w, eigenvalues; V EigenVectors
        w, V = self.SortByEigenValue(w, V)  # after sorting

        # Compute singular values
        sigmas = list(map(lambda x: np.sqrt(x)
                          if x > 0 else 0, w))
        sigmas = np.array(sigmas)
        sigmas_matrix = np.diag(sigmas)

        # Compute the rank of sigmas_matrix
        if Num_Of_sigmas == None:
            rank_of_sigmasMatrix = len(list(filter(lambda x:x>0, sigmas)))
        else:
            rank_of_sigmasMatrix = Num_Of_sigmas
        sigmas_matrix = sigmas_matrix[0:rank_of_sigmasMatrix, :]

        # Compute the left eigenvectors
        U = np.zeros((A.shape[0], rank_of_sigmasMatrix))
        for i in range(rank_of_sigmasMatrix):
            U[:, i] = np.dot(A, V[:, i]) / sigmas[i]

        V = V[:, 0:rank_of_sigmasMatrix]
        sigmas_matrix = sigmas_matrix[0:rank_of_sigmasMatrix, 0: rank_of_sigmasMatrix]

        self.U = U
        self.sigmas = sigmas_matrix
        self.V = V

        return U, sigmas_matrix, V

    def SortByEigenValue(self, Eigenvalues, EigenVectors):
        index = np.argsort((-1*Eigenvalues))
        Eigenvalues = Eigenvalues[index]
        EigenVectors = EigenVectors[:, index]

        return Eigenvalues, EigenVectors

    def RebuildMatrix(self):
        U = self.U
        V = self.V
        sigmas = self.sigmas
        A = np.dot((np.dot(U, sigmas)),V.T)
        return A


if __name__ == '__main__':
    A = np.array([[1,0,0,0],
                  [0,0,1,4],
                  [0,3,0,0],
                  [0,5,0,0],
                  [2,0,0,0]])
    svd = SVD(A)
    svd.SVD(A)
    AA = svd.RebuildMatrix()
    print(AA)
